#! /usr/bin/python
# -*- coding: UTF-8 -*-

import random
import matplotlib.pyplot as plt
import math
from cPickle import load,dump

u = {}
freq = {}
ints = {}
network = 'newcontacts.txt'
start_person = '1112942'
#start_person = '484181'
totaltime = 190
stophold = 3400
temp = open('temp.txt','w')

random.seed()
intsitems = []

def get_wait_time(id):
    if network == 'randcontacts.txt':
        return random.choice(range(1,30))

    fr = freq.get(int(id))
    if fr:
        return 30/fr
    else:
        return random.choice(range(16,30))

def get_interest_p(id,pi = 20000):
    if network == 'randcontacts.txt':
        try:
            n_ints = intsitems[int(id)][1]
        except KeyError:
            n_ints = 0
    else:
        n_ints = ints.get(id, 0)
    inp = n_ints/float(pi)
    temp.write("%s\n"%(n_ints/float(pi)))
    return inp

def init_relation():
    global u,freq,ints,intsitems

    print 'loading freq info...'
    freq = load(open('user_freq.pkl'))[0]

    print 'loading interest info...'
    ints = load(open('interest_dict.pkl'))
    intsitems = ints.items()

    print 'loading network info...'
    f = open(network)

    print 'init relations...'
    for i in f:
        j = i[:-1].split(',')
        try:
            u[j[1]].append(j[0])
        except KeyError:
            u[j[1]] = [j[0]]

def do_spread(sp=start_person,px=0.9,dpx=5,pi=20000,
        p1only=False,p2only=False,rmode=True):
    t = 0
    ipeople = set()
    rpeople = set()
    openat = {}
    ioftime = []
    ipeople.update([sp])
    ioftime.append([sp])
    people_in = [sp]
    print 'start with %s' % sp
    for t in xrange(1,totaltime):
        ioftime.append([])
        if len(people_in) > stophold:
            break
        willsend = set()
        for i in people_in:
            willsend.update(u[i])
        print 'willsend to:%s, including ipeople %s' % (
                len(willsend),len(ipeople))
        if rmode:
            willsend = willsend - ipeople - rpeople
        else:
            willsend = willsend - ipeople
        for s in willsend:
            waittime = t + get_wait_time(s)
            if openat.get(waittime) is None:
                openat[waittime] = [s]
            else:
                openat[waittime].append(s)
        ontime = openat.get(t, [])
        print '%s people receive'%len(ontime)
        people_in = set()
        for s in ontime:
            p1 = 0 if p2only else get_interest_p(s,pi=pi)
            p2 = 0 if p1only else float(t)**(-px)/dpx
            interest_p = p1 + p2
            deci = random.random()
            if  deci < interest_p:
                people_in.update([s])
                ioftime[t].append(s)
            elif rmode:
                rpeople.update([s])
        print '%s people in!'%len(people_in)
        ipeople.update(people_in)
        t += 1
        print ''
    for t,i in enumerate(ioftime):
        #print '%s: %s'%(t,len(i))
        if t>0:
            sumoftime.append(sumoftime[t-1]+len(i))
        else:
            sumoftime = [len(i)]
    return sumoftime

def dopx(nstart,nrange):
    global px
    allres = []
    for i in range(nstart,nrange):
        px = i/2.0
        res = do_spread()
        allres.append(res.pop())
    plt.title('total spread of different px')
    plt.xlabel('px')
    plt.ylabel('total')
    plt.plot([i/2.0 for i in range(nstart,nrange)],allres,'ro-')
    plt.grid(True)
    plt.show()
    dump(allres,open('temp_res.pkl','w'))

def do(times,sp=start_person,px=0.8,dpx=5,pi=20000,
        p1only=False,p2only=False,rmode=True,style='r-'):
    allres = [0] * totaltime
    for i in range(times):
        res = do_spread(sp=sp,px=px,dpx=dpx,pi=pi,
                p2only=p2only,p1only=p1only,rmode=rmode)
        for n,r in enumerate(res):
            allres[n] += r
    averes = [i/float(times) for i in allres]
    plt.title('average spread of %s times'%times)
    plt.xlabel('t')
    plt.ylabel('n')
    plt.plot(averes,style)
    plt.grid(True)
    plt.show()
    dump(averes,open('temp_res.pkl','w'))
    return averes

if __name__ == '__main__':
    init_relation()
    do(1)

